import chisel3._
import chisel3.util.Decoupled

class EXU(xlen: Int) extends Module {
  val io = IO(new Bundle {
    val in = Flipped(Decoupled(new MessageDE(xlen)))
    val out =  Decoupled(new MessageEM(xlen))

    // from pipeline register of MEU
    val br_taken = Input(Bool())
  })

  // pipeline registers (idu => exu)
  val exu_inst     = RegInit(0.U(xlen.W))
  val exu_pc       = RegInit(0.U(xlen.W))
  val exu_PC_sel   = RegInit(0.U(2.W))
  val exu_Imm_sel  = RegInit(0.U(3.W))
  val exu_Alu_op   = RegInit(0.U(4.W))
  val exu_st_type  = RegInit(0.U(2.W))
  val exu_ld_type  = RegInit(0.U(3.W))
  val exu_br_type  = RegInit(0.U(3.W))
  val exu_A_sel    = RegInit(0.U(1.W))
  val exu_B_sel    = RegInit(0.U(1.W))
  val exu_wb_sel   = RegInit(0.U(2.W))
  val exu_csr_cmd  = RegInit(0.U(3.W))
  val exu_wb_en    = RegInit(0.B)
  val exu_wb_addr  = RegInit(0.U(5.W))
  val exu_rs1_data = RegInit(0.U(xlen.W))
  val exu_rs2_data = RegInit(0.U(xlen.W))

  // valid for output
  val valid = RegInit(0.B)

  // ready for input
  val ready = RegInit(1.B)

  val immGen  = Module(new ImmGen(xlen))
  val alu     = Module(new Alu(xlen, xlen))
  val brcond  = Module(new BrCond(xlen))

  exu_inst     := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.inst,     exu_inst))
  exu_pc       := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.pc,       exu_pc))
  exu_PC_sel   := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.PC_sel,   exu_PC_sel))
  exu_Imm_sel  := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.Imm_sel,  exu_Imm_sel))
  exu_Alu_op   := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.Alu_op,   exu_Alu_op))
  exu_st_type  := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.st_type,  exu_st_type))
  exu_ld_type  := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.ld_type,  exu_ld_type))
  exu_br_type  := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.br_type,  exu_br_type))
  exu_A_sel    := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.A_sel,    exu_A_sel))
  exu_B_sel    := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.B_sel,    exu_B_sel))
  exu_wb_sel   := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.wb_sel,   exu_wb_sel))
  exu_csr_cmd  := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.csr_cmd,  exu_csr_cmd))
  exu_wb_en    := Mux(io.br_taken, 0.B, Mux(io.in.valid && io.in.ready, io.in.bits.wb_en,    exu_wb_en))
  exu_wb_addr  := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.wb_addr,  exu_wb_addr))
  exu_rs1_data := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.rs1_data, exu_rs1_data))
  exu_rs2_data := Mux(io.br_taken, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.rs2_data, exu_rs2_data))

  immGen.io.inst    := exu_inst
  immGen.io.Imm_sel := exu_Imm_sel

  alu.io.imm      := immGen.io.imm
  alu.io.rs1_data := exu_rs1_data
  alu.io.rs2_data := exu_rs2_data
  alu.io.pc_out   := exu_pc
  alu.io.Alu_op   := exu_Alu_op
  alu.io.A_sel    := exu_A_sel
  alu.io.B_sel    := exu_B_sel

  brcond.io.br_type  := exu_br_type
  brcond.io.rs1_data := exu_rs1_data
  brcond.io.rs2_data := exu_rs2_data

  // io.br_taken := brcond.io.br_taken | Mux(exu_inst(6, 0) === "b1101111".U, 1.B, Mux(exu_inst(6, 0) === "b1100111".U && exu_inst(14, 12) === "b000".U, 1.B, 0.B))

  io.out.bits.inst    := Mux(io.br_taken, 0.U, exu_inst)
  io.out.bits.pc      := Mux(io.br_taken, 0.U, exu_pc)

  io.out.bits.PC_sel  := Mux(io.br_taken, 0.U, exu_PC_sel)
  io.out.bits.st_type := Mux(io.br_taken, 0.U, exu_st_type)
  io.out.bits.ld_type := Mux(io.br_taken, 0.U, exu_ld_type)
  io.out.bits.wb_sel  := Mux(io.br_taken, 0.U, exu_wb_sel)
  io.out.bits.csr_cmd := Mux(io.br_taken, 0.U, exu_csr_cmd)
  io.out.bits.wb_en   := Mux(io.br_taken, 0.B, exu_wb_en)
  io.out.bits.br_taken:= Mux(io.br_taken, 0.B, brcond.io.br_taken | Mux(exu_inst(6, 0) === "b1101111".U, 1.B, Mux(exu_inst(6, 0) === "b1100111".U && exu_inst(14, 12) === "b000".U, 1.B, 0.B))) // NOTE: for simulation to identify next pc

  io.out.bits.Alu_out  := Mux(io.br_taken, 0.U, alu.io.Alu_out)
  io.out.bits.wb_addr  := Mux(io.br_taken, 0.U, exu_wb_addr)
  io.out.bits.rs2_data := Mux(io.br_taken, 0.U, exu_rs2_data)

  valid := Mux(io.out.valid & io.out.ready, 0.B, Mux(io.in.valid & io.in.ready, 1.B, valid))
  io.out.valid := valid

  ready := Mux(io.in.valid & io.in.ready, 0.B, Mux(io.out.valid & io.out.ready, 1.B, ready))
  io.in.ready  := ready
}
